import random
import shutil
import os  # 用于遍历文件夹
from pathlib import Path
from PIL import Image
import SimpleITK as sitk
import imageio  # 图像io
import nibabel as nib
import numpy as np


def random_select_files(folder_path, target_path, num_files, seed=None):
    """
    从指定文件夹中随机抽取指定数量的文件，并将其复制到当前目录下的'random_files'文件夹中
    :param folder_path: 文件夹路径
    :param num_files: 需要抽取的文件数量
    :param seed: 随机种子
    """
    if not os.path.isdir(folder_path):
        print(f"{folder_path}不是一个有效的文件夹路径")
        return
    files = os.listdir(folder_path)
    if num_files > len(files):
        print("文件数量不足，无法完成抽取")
        return
    if seed:
        random.seed(seed)
    random.shuffle(files)
    random_files = files[:num_files]
    if not os.path.exists(target_path):
        os.mkdir(target_path)
    for file in random_files:
        shutil.copy(os.path.join(folder_path, file), target_path)
    print(f"{num_files}个文件已被成功抽取到{target_path}文件夹中")


def produceImage(file_in, width, height, file_out):
    image = Image.open(file_in)
    resized_image = image.resize((width, height), Image.ANTIALIAS)
    resized_image.save(os.path.join(file_out, str(file_in).split('\\')[2]))
    # resized_image.save(file_out, quality=95, dpi=(300.0, 300.0)) # 调整图像的分辨率为300,dpi可以更改


def convert_to_nii(input_file, output_file):  # 读取输入文件
    image = sitk.ReadImage(input_file)
    # 将图像保存为NIFTI格式
    sitk.WriteImage(image, output_file)
    # 调用函数进行转换


def random_files():
    # 示例
    for num in range(3):
        folder_path = f'./data/data_class/{num + 1}'  # 替换为实际文件夹路径
        target_path = f'./data_classification/{num + 1}'
        num_files = 2500  # 需要抽取的文件数量
        seed = 123  # 随机种子
        random_select_files(folder_path, target_path, num_files, seed)


def nii_process_png(type_id):
    data_route = './data/data_class'
    BASEDIR = os.path.normpath(f'./data/{type_id}')
    files = Path(BASEDIR).rglob(f'*-{type_id}-Quality*')

    for filename in files:
        print(filename)
        class_data_path = os.path.join(data_route, str(filename).split('-')[len(str(filename).split('-')) - 1][0:1])
        if not os.path.isdir(class_data_path):
            os.mkdir(class_data_path)
        imageNii = nib.load(str(filename))
        image = imageNii.get_fdata()

        for i in range(image.shape[-3]):
            # tempResult = image[:, :, i]
            new_image = nib.Nifti1Image(np.array([image[:, :, i, 0, 0]]).transpose((-1, 1, 0)), np.eye(4))
            silce = np.array([image[:, :, i, 0, 0]]).transpose((-1, 1, 0))
            silce = silce.astype(np.uint8)
            imageio.imwrite(
                os.path.join(class_data_path, '{}.png'.format(str(filename).split('\\')[2][:-7] + '_' + str(i))),
                silce)  # 保存图像
            # nib.save(new_image, f'{route}/{file}_{"%04d" % i}_0000.nii.gz')
            print(f'{class_data_path}/{"%04d" % i}.nii.gz')


def convert_224():
    for num in range(3):
        BASEDIR = os.path.normpath(f'./data_classification/{num + 1}')
        file_out = os.path.normpath(f'./data_224/{num + 1}')
        files = Path(BASEDIR).rglob('*.png')
        if not os.path.isdir(file_out):
            os.mkdir(file_out)
        for filename in files:
            print(filename)
            produceImage(filename, 224, 224, file_out)


def file_filtering(type_id):
    BASEDIR = os.path.normpath(r'E:\data')
    files = Path(BASEDIR).rglob(f'*-{type_id}-Quality*')
    variables = {}
    for i in range(3):
        variables[f'num_{i + 1}'] = 0
    for class_id in files:
        name = str(class_id).split('\\')[len(str(class_id).split('\\')) - 1][:-6]
        quality_type = int(name.split('_')[0][-1])
        quality_path = f'data/{type_id}/{name}.nii.gz'
        print(name)
        for i in range(3):
            if quality_type == i + 1 and variables[f'num_{i + 1}'] < 20:
                variables[f'num_{i + 1}'] += 1
                convert_to_nii(str(class_id), quality_path)


# 复制文件路径 抽取指定ID文件
def file_extract(type_id):
    data_file = f'F:/data/{type_id}'
    if not os.path.isdir(data_file):
        os.mkdir(data_file)
    BASEDIR = os.path.normpath(f'F:\李恒\分类数据')
    files = Path(BASEDIR).rglob(f'*-{type_id}-Quality*')
    variables = {}
    for i in range(3):
        variables[f'num_{i + 1}'] = 0
    for class_id in files:
        name = str(class_id).split('\\')[len(str(class_id).split('\\')) - 1][:-6]
        quality_type = int(name.split('_')[0][-1])
        print(name)
        for i in range(3):
            if quality_type == i + 1 and variables[f'num_{i + 1}'] < 40:
                variables[f'num_{i + 1}'] += 1
                shutil.copy(str(class_id), data_file)


def nii_to_image_z(img_path, route, color_base):
    img = nib.load(img_path)
    img_fdata = img.get_fdata()
    (x, y, z) = img.shape  # 获取图像的3个方向的维度

    for i in range(z):  # z方向
        # silce = np.fliplr(np.rot90(img_fdata[:, :, i], -1))
        silce = img_fdata[:, :, i]
        silce = silce.astype(np.uint8)
        # imageio.imwrite(os.path.join(route, '{}.png'.format(img_path.split('\\')[6][:-7])), silce * color_base)  # 保存图像
        imageio.imwrite(os.path.join(route, '{}.png'.format(img_path.split('\\')[6][:-7])), silce)  # 保存图像


# dicom转换nii.gz 36 11 37 21 20————————1 25
type_id = 36
# file_filtering(type_id)
# nii.gz转png
# nii_process_png(type_id)
# 抽取文件
# random_files()
# 转换文件格式224
# convert_224()
#
file_extract(type_id)
